#!/bin/bash
#SBATCH -p kshdtest       # 分区名称
#SBATCH --gres=dcu:2      # 请求1个DCU资源
#SBATCH -n 1            # MPI进程数
#SBATCH -c 16            # CPU核心数 (与 OMP_NUM_THREADS 匹配)
#SBATCH -t 05:00:00     # 最大运行时间
#SBATCH -J OpenBLASMatMulTest # 作业名称 (更新)
#SBATCH -o %x_%j.out    # 标准输出文件 (%x = Job name, %j = Job ID)
#SBATCH -e %x_%j.err    # 错误输出文件

# --- OpenBLAS Performance Test Script ---

echo "=== OpenBLAS Test Starting: $(date) ===="
echo "Job ID: $SLURM_JOB_ID"
echo "Node: $SLURM_JOB_NODELIST"
echo "Working Directory: $(pwd)"
echo "CPUs per task requested: $SLURM_CPUS_PER_TASK"
export OMP_NUM_THREADS=$SLURM_CPUS_PER_TASK
echo "OMP_NUM_THREADS set to: $OMP_NUM_THREADS"

module purge
module load compiler/rocm/dtk-23.10
module load compiler/gcc/9.3.0
module load mathlib/blas/3.5.0/intel

# --- Set OpenBLAS/OMP Threads ---
# Suggest OpenBLAS uses the number of cores allocated by Slurm
export OMP_NUM_THREADS=$SLURM_CPUS_PER_TASK
echo "Setting OMP_NUM_THREADS to: $OMP_NUM_THREADS"

echo "=== Environment ===="
echo "HIP_PATH: ${HIP_PATH}"
echo "ROCM_PATH: ${ROCM_PATH}"
echo ""

# 2. Compile Code (Using modified Makefile)
echo "--- Compiling Code ---"
make clean
EXECUTABLE="./matmul" # Ensure this matches Makefile target
make -j HIPCC=hipcc

if [ $? -ne 0 ]; then
    echo "!!! Compilation FAILED !!!"
    exit 1
fi
echo "Compilation successful."
echo ""

# 3. Run Adaptive Tests
echo "--- Running Adaptive Performance Tests (with OpenBLAS) ---"
LOG_FILE="openblas_test_output.log" # New log file name
SUMMARY_FILE="openblas_summary.log"   # New summary file name
echo "Test output will be logged to $LOG_FILE"
echo "Summary will be logged to $SUMMARY_FILE"
echo "" > $LOG_FILE # Clear log file
echo "Size,Best_Hybrid_Time_ms,Best_CPU_Ratio,Best_Iteration,Avg_CPU_Time_ms,Avg_DCU_Time_ms" > $SUMMARY_FILE # CSV Header

SIZES=(512 1024 2048 4096 8192 16384 32768)
ADAPTIVE_ITERATIONS=10
INITIAL_HYBRID_RATIO=0.5

# --- Test Loop Logic (Identical to quick_test_job.sh) ---
for SIZE in ${SIZES[@]}; do
    M=$SIZE
    K=$SIZE
    N=$SIZE

    echo "--- Testing Size: ${M}x${K}x${N} ---" | tee -a $LOG_FILE

    # --- Reset state ---
    declare -a cpu_times_ms
    declare -a dcu_times_ms
    declare -a hybrid_times_ms
    declare -a hybrid_ratios_used

    cumulative_cpu_time_ms=0.0
    cumulative_dcu_time_ms=0.0
    total_cpu_runs=0
    total_dcu_runs=0

    best_hybrid_avg_time_ms=-1.0
    best_hybrid_ratio=-1.0
    best_hybrid_iteration=-1

    current_cpu_ratio=$INITIAL_HYBRID_RATIO

    # --- Adaptive Iteration Loop ---
    for iter in $(seq 1 $ADAPTIVE_ITERATIONS); do
        echo -e "
--- Size $SIZE - Iteration $iter / $ADAPTIVE_ITERATIONS ---" | tee -a $LOG_FILE

        # --- Run Pure CPU Mode (single run) ---
        echo "Running CPU Mode (1 time)..." | tee -a $LOG_FILE
        output_cpu=$(echo "$M $K $N" | $EXECUTABLE -c)
        time_ms_cpu=$(echo "$output_cpu" | grep "RESULT_TIME_MS:" | awk '{print $2}')
        if [[ -z "$time_ms_cpu" || ! "$time_ms_cpu" =~ ^[0-9]+([.][0-9]+)?$ ]]; then
            echo "错误：无法从 CPU 运行 $iter 的输出中解析有效时间！ Output:" | tee -a $LOG_FILE
            echo "$output_cpu" | tee -a $LOG_FILE
            exit 1 # Exit if parsing fails
        fi
        cpu_times_ms+=($time_ms_cpu)
        cumulative_cpu_time_ms=$(echo "$cumulative_cpu_time_ms + $time_ms_cpu" | bc)
        total_cpu_runs=$(($total_cpu_runs + 1))
        echo "CPU Time (this iter): $time_ms_cpu ms" | tee -a $LOG_FILE

        # --- Run Pure DCU Mode (single run) ---
        echo "Running DCU Mode (1 time)..." | tee -a $LOG_FILE
        output_dcu=$(echo "$M $K $N" | $EXECUTABLE -d)
        time_ms_dcu=$(echo "$output_dcu" | grep "RESULT_TIME_MS:" | awk '{print $2}')
        if [[ -z "$time_ms_dcu" || ! "$time_ms_dcu" =~ ^[0-9]+([.][0-9]+)?$ ]]; then
            echo "错误：无法从 DCU 运行 $iter 的输出中解析有效时间！ Output:" | tee -a $LOG_FILE
            echo "$output_dcu" | tee -a $LOG_FILE
            exit 1 # Exit if parsing fails
        fi
        dcu_times_ms+=($time_ms_dcu)
        cumulative_dcu_time_ms=$(echo "$cumulative_dcu_time_ms + $time_ms_dcu" | bc)
        total_dcu_runs=$(($total_dcu_runs + 1))
        echo "DCU Time (this iter): $time_ms_dcu ms" | tee -a $LOG_FILE

        # --- Calculate cumulative averages and next cpuRatio ---
        if [[ $iter -gt 1 ]]; then
            if [[ $total_cpu_runs -gt 0 && $total_dcu_runs -gt 0 ]]; then
                t_cpu_avg=$(echo "scale=6; $cumulative_cpu_time_ms / $total_cpu_runs" | bc)
                t_dcu_avg=$(echo "scale=6; $cumulative_dcu_time_ms / $total_dcu_runs" | bc)
                echo "Cumulative Avg: CPU=$t_cpu_avg ms, DCU=$t_dcu_avg ms" | tee -a $LOG_FILE

                denominator=$(echo "$t_cpu_avg + $t_dcu_avg" | bc)
                if (( $(echo "$denominator > 0.000001" | bc -l) )); then
                    current_cpu_ratio=$(echo "scale=5; $t_dcu_avg / $denominator" | bc)
                    if (( $(echo "$current_cpu_ratio < 0" | bc -l) )); then current_cpu_ratio=0.0; fi
                    if (( $(echo "$current_cpu_ratio > 1" | bc -l) )); then current_cpu_ratio=1.0; fi
                     echo "Calculated next Hybrid CPU Ratio: $current_cpu_ratio" | tee -a $LOG_FILE
                else
                    echo "Warning: Sum of avg CPU and DCU times is zero or too small. Using default ratio 0.5" | tee -a $LOG_FILE
                    current_cpu_ratio=0.5
                fi
            else
                 echo "Warning: Not enough historical data for CPU/DCU runs. Using previous ratio: $current_cpu_ratio" | tee -a $LOG_FILE
            fi
        else
             echo "Iteration 1, using initial Hybrid CPU Ratio: $current_cpu_ratio" | tee -a $LOG_FILE
        fi
        hybrid_ratios_used+=($current_cpu_ratio)

        # --- Run Hybrid Mode (single run) ---
        echo "Running Hybrid Mode (CPU Ratio: $current_cpu_ratio, 1 time)..." | tee -a $LOG_FILE
        output_hybrid=$(echo "$M $K $N" | $EXECUTABLE -h --force-cpu-ratio $current_cpu_ratio)
        time_ms_hybrid=$(echo "$output_hybrid" | grep "RESULT_TIME_MS:" | awk '{print $2}')
         if [[ -z "$time_ms_hybrid" || ! "$time_ms_hybrid" =~ ^[0-9]+([.][0-9]+)?$ ]]; then
            echo "错误：无法从混合模式运行 $iter 的输出中解析有效时间！ Output:" | tee -a $LOG_FILE
             echo "$output_hybrid" | tee -a $LOG_FILE
            exit 1 # Exit if parsing fails
        fi
        hybrid_times_ms+=($time_ms_hybrid)
        echo "Hybrid Time (this iter): $time_ms_hybrid ms" | tee -a $LOG_FILE

        # --- Update best result ---
        is_first_best=0
        [[ $(echo "$best_hybrid_avg_time_ms == -1.0" | bc -l) -eq 1 ]] && is_first_best=1
        is_better=0
        [[ $is_first_best -eq 0 && $(echo "$time_ms_hybrid < $best_hybrid_avg_time_ms" | bc -l) -eq 1 ]] && is_better=1
        if [[ $is_first_best -eq 1 || $is_better -eq 1 ]]; then
             best_hybrid_avg_time_ms=$time_ms_hybrid
             best_hybrid_ratio=$current_cpu_ratio
             best_hybrid_iteration=$iter
             echo "  *** New best hybrid result for size $SIZE found! (Iter $iter) ***" | tee -a $LOG_FILE
        fi
    done # End adaptive iterations

    # --- Log summary for this size ---
    final_avg_cpu="N/A"
    final_avg_dcu="N/A"
     if [[ $total_cpu_runs -gt 0 ]]; then final_avg_cpu=$(echo "scale=3; $cumulative_cpu_time_ms / $total_cpu_runs" | bc); fi
     if [[ $total_dcu_runs -gt 0 ]]; then final_avg_dcu=$(echo "scale=3; $cumulative_dcu_time_ms / $total_dcu_runs" | bc); fi
    echo "${SIZE}x${SIZE},${best_hybrid_avg_time_ms},${best_hybrid_ratio},${best_hybrid_iteration},${final_avg_cpu},${final_avg_dcu}" >> $SUMMARY_FILE
    echo "--- Finished Size $SIZE ---" | tee -a $LOG_FILE
    echo "" | tee -a $LOG_FILE
done # End SIZES loop
echo "" | tee -a $LOG_FILE
echo "#######################################################" | tee -a $LOG_FILE
echo "--- Running Adaptive Performance Tests (Non-Square Matrices) ---" | tee -a $LOG_FILE
echo "#######################################################" | tee -a $LOG_FILE

SUMMARY_FILE_NONSQUARE="openblas_summary_nonsquare.csv" # New summary file
echo "Summary for NON-SQUARE matrices will be logged to $SUMMARY_FILE_NONSQUARE"
# Use same adaptive iterations and initial ratio, can be changed if needed
# ADAPTIVE_ITERATIONS=10
# INITIAL_HYBRID_RATIO=0.5

# Define non-square dimensions as "MxKxN" strings
NONSQUARE_SIZES=(
    "8192x256x256"   # Extremely Tall-Skinny
    "256x8192x256"   # Extremely Short-Fat
    "4096x2048x1024"  # More General 1
    "1024x4096x2048"  # More General 2
    "65536x2048x2048"  # Extremely Very Tall-Skinny
    "2048x65536x2048"  # Extremely Very Short-Fat
)

# Add header for non-square summary file
echo "MxKxN,Best_Hybrid_Time_ms,Best_CPU_Ratio,Best_Iteration,Avg_CPU_Time_ms,Avg_DCU_Time_ms" > $SUMMARY_FILE_NONSQUARE

# --- Test Loop Logic (Non-Square) ---
for size_str in "${NONSQUARE_SIZES[@]}"; do
    # Parse M, K, N from the string
    IFS='x' read -r M K N <<< "$size_str"

    echo "--- Testing NON-SQUARE Size: ${M}x${K}x${N} ---" | tee -a $LOG_FILE

    # --- Reset state ---
    # ... (declare arrays, reset variables - same as square loop) ...
    declare -a cpu_times_ms
    declare -a dcu_times_ms
    declare -a hybrid_times_ms
    declare -a hybrid_ratios_used
    cumulative_cpu_time_ms=0.0
    cumulative_dcu_time_ms=0.0
    total_cpu_runs=0
    total_dcu_runs=0
    best_hybrid_avg_time_ms=-1.0
    best_hybrid_ratio=-1.0
    best_hybrid_iteration=-1
    current_cpu_ratio=$INITIAL_HYBRID_RATIO

    # --- Adaptive Iteration Loop (Non-Square) ---
    for iter in $(seq 1 $ADAPTIVE_ITERATIONS); do
        echo -e "\n--- NON-SQUARE Size ${M}x${K}x${N} - Iteration $iter / $ADAPTIVE_ITERATIONS ---" | tee -a $LOG_FILE

        # --- Run Pure CPU Mode ---
        # ... (Run CPU, parse time, accumulate - use M, K, N) ...
        echo "Running CPU Mode (1 time)..." | tee -a $LOG_FILE
        output_cpu=$(echo "$M $K $N" | $EXECUTABLE -c) # Use echo pipe
        time_ms_cpu=$(echo "$output_cpu" | grep "RESULT_TIME_MS:" | awk '{print $2}')
        # ... (Error check, accumulate) ...
        if [[ -z "$time_ms_cpu" || ! "$time_ms_cpu" =~ ^[0-9]+([.][0-9]+)?$ ]]; then echo "Error parsing CPU time" | tee -a $LOG_FILE; exit 1; fi
        cpu_times_ms+=($time_ms_cpu)
        cumulative_cpu_time_ms=$(echo "$cumulative_cpu_time_ms + $time_ms_cpu" | bc)
        total_cpu_runs=$(($total_cpu_runs + 1))
        echo "CPU Time (this iter): $time_ms_cpu ms" | tee -a $LOG_FILE

        # --- Run Pure DCU Mode ---
        # ... (Run DCU, parse time, accumulate - use M, K, N) ...
        echo "Running DCU Mode (1 time)..." | tee -a $LOG_FILE
        output_dcu=$(echo "$M $K $N" | $EXECUTABLE -d) # Use echo pipe
        time_ms_dcu=$(echo "$output_dcu" | grep "RESULT_TIME_MS:" | awk '{print $2}')
        # ... (Error check, accumulate) ...
        if [[ -z "$time_ms_dcu" || ! "$time_ms_dcu" =~ ^[0-9]+([.][0-9]+)?$ ]]; then echo "Error parsing DCU time" | tee -a $LOG_FILE; exit 1; fi
        dcu_times_ms+=($time_ms_dcu)
        cumulative_dcu_time_ms=$(echo "$cumulative_dcu_time_ms + $time_ms_dcu" | bc)
        total_dcu_runs=$(($total_dcu_runs + 1))
        echo "DCU Time (this iter): $time_ms_dcu ms" | tee -a $LOG_FILE

        # --- Calculate next cpuRatio ---
        # ... (Calculate current_cpu_ratio - same logic as square) ...
         if [[ $iter -gt 1 && $total_cpu_runs -gt 0 && $total_dcu_runs -gt 0 ]]; then
             t_cpu_avg=$(echo "scale=6; $cumulative_cpu_time_ms / $total_cpu_runs" | bc)
             t_dcu_avg=$(echo "scale=6; $cumulative_dcu_time_ms / $total_dcu_runs" | bc)
             denominator=$(echo "$t_cpu_avg + $t_dcu_avg" | bc)
             if (( $(echo "$denominator > 0.000001" | bc -l) )); then
                 current_cpu_ratio=$(echo "scale=5; $t_dcu_avg / $denominator" | bc)
                 # Clamp ratio [0, 1]
                 if (( $(echo "$current_cpu_ratio < 0" | bc -l) )); then current_cpu_ratio=0.0; fi
                 if (( $(echo "$current_cpu_ratio > 1" | bc -l) )); then current_cpu_ratio=1.0; fi
                  echo "Calculated next Hybrid CPU Ratio: $current_cpu_ratio" | tee -a $LOG_FILE
             else
                 current_cpu_ratio=0.5 # Fallback ratio
                 echo "Warning: Sum of avg times too small. Using fallback ratio $current_cpu_ratio" | tee -a $LOG_FILE
             fi
        else
             echo "Using initial/previous Hybrid CPU Ratio: $current_cpu_ratio" | tee -a $LOG_FILE
        fi
        hybrid_ratios_used+=($current_cpu_ratio)

        # --- Run Hybrid Mode ---
        # ... (Run Hybrid with current_cpu_ratio, parse time - use M, K, N) ...
        echo "Running Hybrid Mode (CPU Ratio: $current_cpu_ratio, 1 time)..." | tee -a $LOG_FILE
        output_hybrid=$(echo "$M $K $N" | $EXECUTABLE -h --force-cpu-ratio $current_cpu_ratio) # Use echo pipe
        time_ms_hybrid=$(echo "$output_hybrid" | grep "RESULT_TIME_MS:" | awk '{print $2}')
         # ... (Error check) ...
         if [[ -z "$time_ms_hybrid" || ! "$time_ms_hybrid" =~ ^[0-9]+([.][0-9]+)?$ ]]; then echo "Error parsing Hybrid time" | tee -a $LOG_FILE; exit 1; fi
        hybrid_times_ms+=($time_ms_hybrid)
        echo "Hybrid Time (this iter): $time_ms_hybrid ms" | tee -a $LOG_FILE


        # --- Update best result ---
        # ... (Update best_hybrid_avg_time_ms, best_hybrid_ratio, best_hybrid_iteration) ...
        is_first_best=0; [[ $(echo "$best_hybrid_avg_time_ms == -1.0" | bc -l) -eq 1 ]] && is_first_best=1
        is_better=0; [[ $is_first_best -eq 0 && $(echo "$time_ms_hybrid < $best_hybrid_avg_time_ms" | bc -l) -eq 1 ]] && is_better=1
        if [[ $is_first_best -eq 1 || $is_better -eq 1 ]]; then
             best_hybrid_avg_time_ms=$time_ms_hybrid
             best_hybrid_ratio=$current_cpu_ratio
             best_hybrid_iteration=$iter
             echo "  *** New best hybrid result for NON-SQUARE size ${M}x${K}x${N} found! (Iter $iter) ***" | tee -a $LOG_FILE
        fi

    done # End adaptive iterations for non-square size

    # --- Log summary for non-square size ---
    # ... (Calculate final averages, echo to SUMMARY_FILE_NONSQUARE) ...
    final_avg_cpu="N/A"; if [[ $total_cpu_runs -gt 0 ]]; then final_avg_cpu=$(echo "scale=3; $cumulative_cpu_time_ms / $total_cpu_runs" | bc); fi
    final_avg_dcu="N/A"; if [[ $total_dcu_runs -gt 0 ]]; then final_avg_dcu=$(echo "scale=3; $cumulative_dcu_time_ms / $total_dcu_runs" | bc); fi
    echo "${M}x${K}x${N},${best_hybrid_avg_time_ms},${best_hybrid_ratio},${best_hybrid_iteration},${final_avg_cpu},${final_avg_dcu}" >> $SUMMARY_FILE_NONSQUARE
    echo "--- Finished NON-SQUARE Size ${M}x${K}x${N} ---" | tee -a $LOG_FILE
    echo "" | tee -a $LOG_FILE

done # End of NON-SQUARE SIZES loop

# 4. Test Python Environment (Optional plotting)
echo "--- Testing Python Environment ---" | tee -a $LOG_FILE
python3 plot_performance.py --summary_file $SUMMARY_FILE --output_prefix openblas_ >> $LOG_FILE 2>&1
if [ $? -eq 0 ]; then
    echo "Python script execution attempt finished (check log for success/errors)." | tee -a $LOG_FILE
else
    echo "!!! Python script execution FAILED (python3 or modules not found) !!!" | tee -a $LOG_FILE
fi
echo ""

# 5. Finish
echo "=== OpenBLAS Test Finished: $(date) ====" | tee -a $LOG_FILE
echo "Check $LOG_FILE for full test output."
echo "Check $SUMMARY_FILE for best hybrid results per size."
echo "Check Slurm output files (%x_%j.out, %x_%j.err) for other messages."

exit 0